# -*- coding: utf-8 -*-
import tritonclient.http as httpclient
import tritonclient.grpc as grpcclient

class TritonServerClient:
    def __init__(self, ip_addr, http_port, grpc_port, verbose=False):
        http_url = ip_addr + ":" + http_port
        self.http_client = httpclient.InferenceServerClient(url=http_url, verbose=verbose, network_timeout=3)
        grpc_url = ip_addr + ":" + grpc_port
        self.grpc_client = grpcclient.InferenceServerClient(url=grpc_url, verbose=verbose)

    def is_server_alive(self, headers=None, query_params=None, use_http=False):
        try:
            if use_http:
                return self.http_client.is_server_live(headers=headers, query_params=query_params)
            else:
                return self.grpc_client.is_server_live(headers=headers)
        except Exception as e:
            return str(e)

    def is_server_ready(self, headers=None, query_params=None, use_http=False):
        try:
            if use_http:
                return self.http_client.is_server_ready(headers=headers, query_params=query_params)
            else:
                return self.grpc_client.is_server_ready(headers=headers)
        except Exception as e:
            return str(e)

    def is_model_ready(self, model_name, model_version="", headers=None, query_params=None, use_http=False):
        try:
            if use_http:
                return self.http_client.is_model_ready(model_name=model_name, model_version=model_version, 
                    headers=headers, query_params=query_params)
            else:
                return self.grpc_client.is_model_ready(model_name=model_name, model_version=model_version, 
                    headers=headers)
        except Exception as e:
            return str(e)

    def get_server_metadata(self, headers=None, query_params=None, use_http=False):
        try:
            if use_http:
                return self.http_client.get_server_metadata(headers=headers, query_params=query_params)
            else:
                return self.grpc_client.get_server_metadata(headers=headers, as_json=True)
        except Exception as e:
            return str(e)

    def get_model_metadata(self, model_name, model_version="", headers=None, query_params=None, use_http=False):
        try:
            if use_http:
                return self.http_client.get_model_metadata(model_name=model_name, model_version=model_version, 
                    headers=headers, query_params=query_params)
            else:
                return self.grpc_client.get_model_metadata(model_name=model_name, model_version=model_version, 
                    headers=headers, as_json=True)
        except Exception as e:
            return str(e)

    def get_model_config(self, model_name, model_version="", headers=None, query_params=None, use_http=False):
        try:
            if use_http:
                return self.http_client.get_model_config(model_name=model_name, model_version=model_version, 
                    headers=headers, query_params=query_params)
            else:
                return self.grpc_client.get_model_config(model_name=model_name, model_version=model_version, 
                    headers=headers, as_json=True)
        except Exception as e:
            return str(e)

    def load_model(self, model_name, headers=None, query_params=None, config=None, files=None, use_http=False):
        try:
            if use_http:
                return self.http_client.load_model(model_name=model_name, headers=headers, 
                    query_params=query_params, config=config, files=files)
            else:
                return self.grpc_client.load_model(model_name=model_name, headers=headers, 
                    config=config, files=files)
        except Exception as e:
            return str(e)

    def unload_model(self, model_name, headers=None, query_params=None, unload_dependents=False, use_http=False):
        try:
            if use_http:
                return self.http_client.unload_model(model_name=model_name, headers=headers, 
                    query_params=query_params, unload_dependents=unload_dependents)
            else:
                return self.grpc_client.unload_model(model_name=model_name, headers=headers, 
                    unload_dependents=unload_dependents)
        except Exception as e:
            return str(e)

    def get_inference_statistics(self, model_name="", model_version="", headers=None, query_params=None, use_http=False):
        try:
            if use_http:
                return self.http_client.get_inference_statistics(model_name=model_name, model_version=model_version, 
                    headers=headers, query_params=query_params)
            else:
                return self.grpc_client.get_inference_statistics(model_name=model_name, model_version=model_version, 
                    headers=headers, as_json=True)
        except Exception as e:
            return str(e)

    